package main

import (
	"fmt"
	"image/png"
	"os"
	"xmediaEmu/pkg/decoder"
	"xmediaEmu/pkg/decoder/codec/h264parser"
	"xmediaEmu/pkg/media/h264reader"
)

// h264 to image.
func main() {
	// 初始化.
	// 读取文件.
	var err error
	videoFile, err := os.Open("h5video.h264")
	if err != nil {
		panic(err)
	}
	reader, err := h264reader.NewReader(videoFile)
	if err != nil {
		panic(err)
	}

	nal, err := reader.NextNAL()
	if err != nil {
		panic(err)
	}

	var nalSps *h264reader.NAL
	var nalPps *h264reader.NAL
	var codecdata h264parser.CodecData
	var videoDec *decoder.VideoDecoder
	isInited := false
	counter := 0
	// var tmpBuff []byte
	for ; nal != nil && err == nil; nal, err = reader.NextNAL() {
		if nal.UnitType == h264reader.NalUnitTypeSPS {
			nalSps = nal
		} else if nal.UnitType == h264reader.NalUnitTypePPS {
			nalPps = nal
		}

		if nalSps != nil && nalPps != nil && !isInited {
			fmt.Println("h264parser inited.")
			codecdata, err = h264parser.NewCodecDataFromSPSAndPPS(nalSps.Data, nalPps.Data)
			isInited = true

			videoDec, err = decoder.NewVideoDecoder(codecdata)
			if err != nil {
				panic(err)
			}
		}

		// 前后关联解析.
		// 下面代码参考:https://github.com/deepch/decode-encode-h264-demo/blob/master/main.go.
		if isInited && (nal.UnitType == h264reader.NalUnitTypeCodedSliceNonIdr || nal.UnitType == h264reader.NalUnitTypeCodedSliceIdr) {
			tmpBuff := append([]byte("\000\000\000\001"), nal.Data...)
			// tmpBuff := append(tmpBuff, nal.Data...)
			frame, err := videoDec.Decode(tmpBuff)
			if err != nil {
				fmt.Printf("Decode failed, data lenght:%d , type:%d \n", len(tmpBuff), nal.UnitType)
				fmt.Println(err)
				continue
			}

			if frame == nil {
				fmt.Printf("frame == nil, ignored data lenght:%d , type:%d \n", len(tmpBuff), nal.UnitType)
				continue
			}

			tmpBuff = tmpBuff[:0]
			filename := fmt.Sprintf("test_%d.png", counter)
			file, err := os.Create(filename)
			if err != nil {
				fmt.Println(err)
				panic(err)
			}

			// YCbCr to image.
			err = png.Encode(file, &frame.Image)
			if err != nil {
				_ = os.Remove(filename)
				fmt.Println(err)
				continue
			}
			counter++
		} else {
			fmt.Printf("Not valid video frames , type:%d \n", nal.UnitType)
		}
	}

	// 另外一种尝试.
	fmt.Printf("All frames parsed!, count:%d \n", counter)
}
